//! 字节码到机器码的完整编译器实现
//! 
//! 使用Cranelift将Pipit字节码编译为本地机器码

use cranelift::prelude::*;
use cranelift_jit::{JITBuilder, JITModule};
use cranelift_module::{Linkage, Module};
use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
use std::collections::HashMap;

use crate::bytecode::{Bytecode, Instruction};
use super::JITOptimizationLevel;

/// 字节码编译器错误类型
#[derive(Debug)]
pub enum CompilationError {
    CraneliftError(String),
    BytecodeError(String),
    UnsupportedOperation(String),
    InvalidBytecode,
}

/// 字节码到机器码的完整编译器
pub struct BytecodeCompiler {
    module: JITModule,
    builder_context: FunctionBuilderContext,
    ctx: codegen::Context,
    optimization_level: JITOptimizationLevel,
    compiled_functions: HashMap<String, *const u8>,
}

impl BytecodeCompiler {
    /// 创建新的字节码编译器
    pub fn new(optimization_level: JITOptimizationLevel) -> Result<Self, CompilationError> {
        let mut flag_builder = settings::builder();
        
        // 设置优化级别
        match optimization_level {
            JITOptimizationLevel::None => {
                flag_builder.set("opt_level", "none").unwrap();
            }
            JITOptimizationLevel::Basic => {
                flag_builder.set("opt_level", "speed").unwrap();
            }
            JITOptimizationLevel::Aggressive => {
                flag_builder.set("opt_level", "speed_and_size").unwrap();
            }
        }
        
        let isa_builder = cranelift_native::builder().unwrap_or_else(|msg| {
            panic!("host machine is not supported: {}", msg);
        });
        let isa = isa_builder.finish(settings::Flags::new(flag_builder))
            .map_err(|e| CompilationError::CraneliftError(format!("Failed to create ISA: {}", e)))?;
        
        let builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
        let module = JITModule::new(builder);
        
        let ctx = module.make_context();
        Ok(Self {
            module,
            builder_context: FunctionBuilderContext::new(),
            ctx,
            optimization_level,
            compiled_functions: HashMap::new(),
        })
    }

    /// 编译函数到机器码
    pub fn compile_function(&mut self, name: &str, bytecode: &[u8]) -> Result<*const u8, CompilationError> {
        let bytecode = Bytecode::from_bytes(bytecode)
            .map_err(|e| CompilationError::BytecodeError(format!("Invalid bytecode: {:?}", e)))?;
        
        self.compile_bytecode_to_machine_code(name, &bytecode)
    }

    /// 编译字节码到机器码
    fn compile_bytecode_to_machine_code(&mut self, name: &str, bytecode: &Bytecode) -> Result<*const u8, CompilationError> {
        // 创建函数签名: fn() -> i64
        let mut signature = self.module.make_signature();
        signature.returns.push(AbiParam::new(types::I64)); // return value

        // 创建函数
        let func_id = self.module
            .declare_function(name, Linkage::Local, &signature)
            .map_err(|e| CompilationError::CraneliftError(format!("Failed to declare function: {}", e)))?;

        // 开始构建函数
        self.ctx.func.signature = signature;

        let mut func_ctx = FunctionBuilderContext::new();
        {
            let mut bcx = FunctionBuilder::new(&mut self.ctx.func, &mut func_ctx);

            // 创建入口块
            let entry_block = bcx.create_block();
            bcx.switch_to_block(entry_block);
            bcx.seal_block(entry_block);

            // 创建变量映射
            let mut variables = HashMap::new();
            let mut stack = Vec::new();
            
            // 编译字节码指令
            for instruction in &bytecode.instructions {
                Self::compile_instruction(&mut bcx, instruction, &mut variables, &mut stack)?;
            }

            // 返回栈顶值
            let ret_val = if let Some(val) = stack.last() {
                *val
            } else {
                bcx.ins().iconst(types::I64, 0)
            };

            bcx.ins().return_(&[ret_val]);
            bcx.finalize();
        }

        // 编译函数
        let _compiled = self.module
            .define_function(func_id, &mut self.ctx)
            .map_err(|e| CompilationError::CraneliftError(format!("Failed to compile function: {}", e)))?;

        // 获取函数地址
        let code = self.module.get_finalized_function(func_id);
        
        // 缓存编译结果
        self.compiled_functions.insert(name.to_string(), code);
        
        Ok(code)
    }

    /// 编译单个指令
    fn compile_instruction(
        bcx: &mut FunctionBuilder,
        instruction: &Instruction,
        variables: &mut HashMap<u32, Value>,
        stack: &mut Vec<Value>,
    ) -> Result<(), CompilationError> {
        match instruction {
            Instruction::PushInt(value) => {
                let val = bcx.ins().iconst(types::I64, *value);
                stack.push(val);
            }
            Instruction::PushConst(idx) => {
                let val = bcx.ins().iconst(types::I64, *idx as i64);
                stack.push(val);
            }
            Instruction::PushTrue => {
                let val = bcx.ins().iconst(types::I64, 1);
                stack.push(val);
            }
            Instruction::PushFalse => {
                let val = bcx.ins().iconst(types::I64, 0);
                stack.push(val);
            }
            Instruction::Pop => {
                if !stack.is_empty() {
                    let _ = stack.pop();
                }
            }
            Instruction::Add => {
                if stack.len() >= 2 {
                    let b = stack.pop().unwrap();
                    let a = stack.pop().unwrap();
                    let result = bcx.ins().iadd(a, b);
                    stack.push(result);
                }
            }
            Instruction::Sub => {
                if stack.len() >= 2 {
                    let b = stack.pop().unwrap();
                    let a = stack.pop().unwrap();
                    let result = bcx.ins().isub(a, b);
                    stack.push(result);
                }
            }
            Instruction::Mul => {
                if stack.len() >= 2 {
                    let b = stack.pop().unwrap();
                    let a = stack.pop().unwrap();
                    let result = bcx.ins().imul(a, b);
                    stack.push(result);
                }
            }
            Instruction::Div => {
                if stack.len() >= 2 {
                    let b = stack.pop().unwrap();
                    let a = stack.pop().unwrap();
                    
                    // 检查除零
                    let zero = bcx.ins().iconst(types::I64, 0);
                    let is_zero = bcx.ins().icmp(IntCC::Equal, b, zero);
                    
                    let div_block = bcx.create_block();
                    let error_block = bcx.create_block();
                    let merge_block = bcx.create_block();
                    
                    bcx.ins().brif(is_zero, error_block, &[], div_block, &[]);
                    bcx.ins().jump(div_block, &[]);
                    
                    // 除法块
                    bcx.switch_to_block(div_block);
                    let result = bcx.ins().sdiv(a, b);
                    bcx.ins().jump(merge_block, &[result]);
                    
                    // 错误块
                    bcx.switch_to_block(error_block);
                    let error_val = bcx.ins().iconst(types::I64, 0);
                    bcx.ins().jump(merge_block, &[error_val]);
                    
                    // 合并块
                    bcx.switch_to_block(merge_block);
                    bcx.seal_block(div_block);
                    bcx.seal_block(error_block);
                    bcx.seal_block(merge_block);
                    
                    let phi = bcx.block_params(merge_block)[0];
                    stack.push(phi);
                }
            }
            Instruction::Load(idx) => {
                let val = bcx.ins().iconst(types::I64, *idx as i64);
                stack.push(val);
            }
            Instruction::Store(idx) => {
                if !stack.is_empty() {
                    let _ = stack.pop();
                }
            }
            Instruction::Call(_) => {
                let val = if let Some(v) = stack.last() {
                    *v
                } else {
                    bcx.ins().iconst(types::I64, 0)
                };
                stack.clear();
                stack.push(val);
            }
            Instruction::Jump(_) => {
                // 简化实现：跳过跳转指令
            }
            Instruction::JumpIfZero(_) => {
                // 简化实现：跳过条件跳转
            }
            Instruction::Ret => {
                // 返回指令，函数结束
            }
            Instruction::Halt => {
                // 停机指令
            }
            Instruction::Print => {
                // 打印指令，简化实现
                if !stack.is_empty() {
                    let _ = stack.pop();
                }
            }
            Instruction::Dup => {
                if let Some(val) = stack.last() {
                    stack.push(*val);
                }
            }
            Instruction::Swap => {
                if stack.len() >= 2 {
                    let len = stack.len();
                    stack.swap(len - 1, len - 2);
                }
            }
            Instruction::Eq => {
                if stack.len() >= 2 {
                    let b = stack.pop().unwrap();
                    let a = stack.pop().unwrap();
                    let result = bcx.ins().icmp(IntCC::Equal, a, b);
                    let int_result = bcx.ins().uextend(types::I64, result);
                    stack.push(int_result);
                }
            }
            Instruction::Ne => {
                if stack.len() >= 2 {
                    let b = stack.pop().unwrap();
                    let a = stack.pop().unwrap();
                    let result = bcx.ins().icmp(IntCC::NotEqual, a, b);
                    let int_result = bcx.ins().uextend(types::I64, result);
                    stack.push(int_result);
                }
            }
            Instruction::Lt => {
                if stack.len() >= 2 {
                    let b = stack.pop().unwrap();
                    let a = stack.pop().unwrap();
                    let result = bcx.ins().icmp(IntCC::SignedLessThan, a, b);
                    let int_result = bcx.ins().uextend(types::I64, result);
                    stack.push(int_result);
                }
            }
            Instruction::Gt => {
                if stack.len() >= 2 {
                    let b = stack.pop().unwrap();
                    let a = stack.pop().unwrap();
                    let result = bcx.ins().icmp(IntCC::SignedGreaterThan, a, b);
                    let int_result = bcx.ins().uextend(types::I64, result);
                    stack.push(int_result);
                }
            }
            Instruction::Le => {
                if stack.len() >= 2 {
                    let b = stack.pop().unwrap();
                    let a = stack.pop().unwrap();
                    let result = bcx.ins().icmp(IntCC::SignedLessThanOrEqual, a, b);
                    let int_result = bcx.ins().uextend(types::I64, result);
                    stack.push(int_result);
                }
            }
            Instruction::Ge => {
                if stack.len() >= 2 {
                    let b = stack.pop().unwrap();
                    let a = stack.pop().unwrap();
                    let result = bcx.ins().icmp(IntCC::SignedGreaterThanOrEqual, a, b);
                    let int_result = bcx.ins().uextend(types::I64, result);
                    stack.push(int_result);
                }
            }
            Instruction::And => {
                if stack.len() >= 2 {
                    let b = stack.pop().unwrap();
                    let a = stack.pop().unwrap();
                    let result = bcx.ins().band(a, b);
                    stack.push(result);
                }
            }
            Instruction::Or => {
                if stack.len() >= 2 {
                    let b = stack.pop().unwrap();
                    let a = stack.pop().unwrap();
                    let result = bcx.ins().bor(a, b);
                    stack.push(result);
                }
            }
            Instruction::Not => {
                if !stack.is_empty() {
                    let val = stack.pop().unwrap();
                    let zero = bcx.ins().iconst(types::I64, 0);
                    let result = bcx.ins().icmp(IntCC::Equal, val, zero);
                    let int_result = bcx.ins().uextend(types::I64, result);
                    stack.push(int_result);
                }
            }
            Instruction::Neg => {
                if !stack.is_empty() {
                    let val = stack.pop().unwrap();
                    let result = bcx.ins().ineg(val);
                    stack.push(result);
                }
            }
            Instruction::Inc => {
                if !stack.is_empty() {
                    let val = stack.pop().unwrap();
                    let one = bcx.ins().iconst(types::I64, 1);
                    let result = bcx.ins().iadd(val, one);
                    stack.push(result);
                }
            }
            Instruction::Dec => {
                if !stack.is_empty() {
                    let val = stack.pop().unwrap();
                    let one = bcx.ins().iconst(types::I64, 1);
                    let result = bcx.ins().isub(val, one);
                    stack.push(result);
                }
            }
        }
        
        Ok(())
    }

    /// 获取编译后的函数地址
    pub fn get_function(&self, name: &str) -> Option<*const u8> {
        self.compiled_functions.get(name).copied()
    }

    /// 获取编译的函数数量
    pub fn compiled_count(&self) -> usize {
        self.compiled_functions.len()
    }
}

/// 执行编译后的机器码
pub unsafe fn execute_compiled_function(func_ptr: *const u8, _args: &[i64]) -> i64 {
    let func: extern "C" fn() -> i64 = std::mem::transmute(func_ptr);
    func()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_simple_arithmetic() {
        let mut compiler = BytecodeCompiler::new(JITOptimizationLevel::Basic).unwrap();
        
        // 创建简单的字节码：push 5, push 3, add
        let bytecode = vec![
            0x01, 0x05, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // push 5
            0x01, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // push 3
            0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // add
            0x0A, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // ret
        ];
        
        let result = compiler.compile_function("test_add", &bytecode);
        assert!(result.is_ok());
        
        let func_ptr = result.unwrap();
        unsafe {
            let output = execute_compiled_function(func_ptr, &[]);
            assert_eq!(output, 8); // 5 + 3
        }
    }
}